package cn.sciento.boot.api.customize.compile;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import javax.tools.Diagnostic;
import javax.tools.DiagnosticCollector;
import javax.tools.JavaCompiler;
import javax.tools.JavaFileManager;
import javax.tools.JavaFileObject;
import javax.tools.StandardJavaFileManager;
import javax.tools.ToolProvider;

public class DynamicCompiler {
  private final JavaCompiler javaCompiler = ToolProvider.getSystemJavaCompiler();
  
  private final StandardJavaFileManager standardFileManager;
  
  private final List<String> options = new ArrayList<>();
  
  private final DynamicClassLoader dynamicClassLoader;
  
  private final Collection<JavaFileObject> compilationUnits = new ArrayList<>();
  
  private final List<Diagnostic<? extends JavaFileObject>> errors = new ArrayList<>();
  
  private final List<Diagnostic<? extends JavaFileObject>> warnings = new ArrayList<>();
  
  public DynamicCompiler(ClassLoader classLoader) {
    if (this.javaCompiler == null)
      throw new IllegalStateException("Can not load JavaCompiler from javax.tools.ToolProvider#getSystemJavaCompiler(), please confirm the application running in JDK not JRE."); 
    this.standardFileManager = this.javaCompiler.getStandardFileManager(null, null, null);
    this.options.add("-Xlint:unchecked");
    this.dynamicClassLoader = new DynamicClassLoader(classLoader);
  }
  
  public void addSource(String className, String source) {
    addSource(new StringSource(className, source));
  }
  
  public void addSource(JavaFileObject javaFileObject) {
    this.compilationUnits.add(javaFileObject);
  }
  
  public Map<String, Class<?>> build() {
    this.errors.clear();
    this.warnings.clear();
    JavaFileManager fileManager = new DynamicJavaFileManager(this.standardFileManager, this.dynamicClassLoader);
    DiagnosticCollector<JavaFileObject> collector = new DiagnosticCollector<>();
    JavaCompiler.CompilationTask task = this.javaCompiler.getTask(null, fileManager, collector, this.options, null, this.compilationUnits);
    try {
      if (!this.compilationUnits.isEmpty()) {
        boolean result = task.call().booleanValue();
        if (!result || collector.getDiagnostics().size() > 0) {
          for (Diagnostic<? extends JavaFileObject> diagnostic : collector.getDiagnostics()) {
            switch (diagnostic.getKind()) {
              case NOTE:
              case MANDATORY_WARNING:
              case WARNING:
                this.warnings.add(diagnostic);
                continue;
            } 
            this.errors.add(diagnostic);
          } 
          if (!this.errors.isEmpty())
            throw new DynamicCompilerException("Compilation Error", this.errors); 
        } 
      } 
      return this.dynamicClassLoader.getClasses();
    } catch (Throwable e) {
      throw new DynamicCompilerException(e, this.errors);
    } finally {
      this.compilationUnits.clear();
    } 
  }
  
  public Map<String, byte[]> buildByteCodes() {
    this.errors.clear();
    this.warnings.clear();
    JavaFileManager fileManager = new DynamicJavaFileManager(this.standardFileManager, this.dynamicClassLoader);
    DiagnosticCollector<JavaFileObject> collector = new DiagnosticCollector<>();
    JavaCompiler.CompilationTask task = this.javaCompiler.getTask(null, fileManager, collector, this.options, null, this.compilationUnits);
    try {
      if (!this.compilationUnits.isEmpty()) {
        boolean result = task.call().booleanValue();
        if (!result || collector.getDiagnostics().size() > 0) {
          for (Diagnostic<? extends JavaFileObject> diagnostic : collector.getDiagnostics()) {
            switch (diagnostic.getKind()) {
              case NOTE:
              case MANDATORY_WARNING:
              case WARNING:
                this.warnings.add(diagnostic);
                continue;
            } 
            this.errors.add(diagnostic);
          } 
          if (!this.errors.isEmpty())
            throw new DynamicCompilerException("Compilation Error", this.errors); 
        } 
      } 
      return this.dynamicClassLoader.getByteCodes();
    } catch (ClassFormatError e) {
      throw new DynamicCompilerException(e, this.errors);
    } finally {
      this.compilationUnits.clear();
    } 
  }
  
  private List<String> diagnosticToString(List<Diagnostic<? extends JavaFileObject>> diagnostics) {
    List<String> diagnosticMessages = new ArrayList<>();
    for (Diagnostic<? extends JavaFileObject> diagnostic : diagnostics)
      diagnosticMessages.add("line: " + diagnostic
          .getLineNumber() + ", message: " + diagnostic.getMessage(Locale.US)); 
    return diagnosticMessages;
  }
  
  public List<String> getErrors() {
    return diagnosticToString(this.errors);
  }
  
  public List<String> getWarnings() {
    return diagnosticToString(this.warnings);
  }
  
  public ClassLoader getClassLoader() {
    return this.dynamicClassLoader;
  }
}
